from typing import Dict, Any
from vnpy.trader.object import BarData
from vnpy.trader.utility import BarGenerator
from vnpy.trader.constant import Interval
from ..template import StrategyTemplate

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,  # 确保正确导入 ArrayManager
)

class AtrRsiStrategy(CtaTemplate):
    """ATR RSI策略"""
    
    author = "用Python的交易员"
    
    # 策略变量
    variables = [
        "atr_value",
        "rsi_value",
        "entry_price"
    ]
    
    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager()
        
        self.buy_signal = False
        self.sell_signal = False
        self.correct_signals = 0  # 正确信号数
        self.total_signals = 0    # 总信号数
        self.last_signal_price = 0.0  # 上一个信号生成时的价格
        
        # 添加策略必需参数
        self.rsi_window = 14
        self.rsi_buy = 40
        self.rsi_sell = 60
        self.atr_window = 20
        self.atr_multiple = 1.0
        self.fixed_size = 1  # 需要定义固定手数
        self.trailing_percent = 0.5  # 需要定义跟踪百分比
        
    def on_init(self):
        """初始化策略"""
        self.write_log("策略初始化")
        
    def on_start(self):
        """启动策略"""
        self.write_log("策略启动")
        
    def on_stop(self):
        """停止策略"""
        self.write_log("策略停止")
        
    # def on_bar(self, bar: BarData):
    #     """
    #     Callback of new bar data update.
    #     """
    #     am = self.am
    #     am.update_bar(bar)
        # if not am.inited:
        #     print(f"ArrayManager inited: {am.inited}, Bar count: {len(am.close)}")  # 添加日志输出
        #     return

    #     atr_array = am.atr(self.atr_window, array=True)
    #     self.atr_value = atr_array[-1]
    #     self.atr_ma = atr_array[-self.atr_window:].mean()
    #     self.rsi_value = am.rsi(self.rsi_window)

    #     self.buy_signal = False
    #     self.sell_signal = False

    #     if self.pos == 0:
    #         self.intra_trade_high = bar.high_price
    #         self.intra_trade_low = bar.low_price

    #         if self.atr_value > self.atr_ma:
    #             if self.rsi_value > self.rsi_buy:
    #                 self.buy(bar.close_price + 5, self.fixed_size)
    #                 self.buy_signal = True  # 更新买入信号
    #                 self.total_signals += 1  # 增加总信号数
    #                 self.last_signal_price = bar.close_price  # 记录信号价格
    #             elif self.rsi_value < self.rsi_sell:
    #                 self.short(bar.close_price - 5, self.fixed_size)
    #                 self.sell_signal = True  # 更新卖出信号
    #                 self.total_signals += 1  # 增加总信号数
    #                 self.last_signal_price = bar.close_price  # 记录信号价格

    #     elif self.pos > 0:
    #         self.intra_trade_high = max(self.intra_trade_high, bar.high_price)
    #         self.intra_trade_low = bar.low_price

    #         long_stop = self.intra_trade_high * \
    #             (1 - self.trailing_percent / 100)
    #         self.sell(long_stop, abs(self.pos), stop=True)

    #     elif self.pos < 0:
    #         self.intra_trade_low = min(self.intra_trade_low, bar.low_price)
    #         self.intra_trade_high = bar.high_price

    #         short_stop = self.intra_trade_low * \
    #             (1 + self.trailing_percent / 100)
    #         self.cover(short_stop, abs(self.pos), stop=True)

    #     self.put_event()
    def on_bar(self, bar: BarData):
        """
        Callback of new bar data update.
        """
        am = self.am
        am.update_bar(bar)
        
        # if not am.inited:
        #     print(f"ArrayManager inited: {am.inited}, Bar count: {len(am.close)}")  # 添加日志输出
        #     return

        # 计算 ATR 和 RSI
        atr_array = am.atr(self.atr_window, array=True)
        self.atr_value = atr_array[-1]
        self.atr_ma = atr_array[-self.atr_window:].mean()
        self.rsi_value = am.rsi(self.rsi_window)

        # 初始化信号状态
        self.buy_signal = False
        self.sell_signal = False

        if self.pos == 0:
            # 平仓状态
            self.intra_trade_high = bar.high_price
            self.intra_trade_low = bar.low_price

            if self.atr_value > self.atr_ma:
                if self.rsi_value > self.rsi_buy:
                    # 生成买入信号
                    self.buy(bar.close_price + 5, self.fixed_size)
                    self.buy_signal = True
                    self.total_signals += 1
                    self.last_signal_price = bar.close_price
                    print(f"生成买入信号，价格: {bar.close_price:.2f}")
                elif self.rsi_value < self.rsi_sell:
                    # 生成卖出信号
                    self.short(bar.close_price - 5, self.fixed_size)
                    self.sell_signal = True
                    self.total_signals += 1
                    self.last_signal_price = bar.close_price
                    print(f"生成卖出信号，价格: {bar.close_price:.2f}")

        elif self.pos > 0:
            # 多头持仓状态
            self.intra_trade_high = max(self.intra_trade_high, bar.high_price)
            self.intra_trade_low = bar.low_price

            # 检查买入信号是否正确
            if self.last_signal_price > 0 and bar.close_price > self.last_signal_price:
                self.correct_signals += 1
                print(f"买入信号正确: 价格从 {self.last_signal_price:.2f} 上涨到 {bar.close_price:.2f}")
                self.last_signal_price = 0.0

            # 设置止损点
            long_stop = self.intra_trade_high * (1 - self.trailing_percent / 100)
            self.sell(long_stop, abs(self.pos), stop=True)
            print(f"多头持仓，当前价格: {bar.close_price:.2f}, 止损价格: {long_stop:.2f}")

        elif self.pos < 0:
            # 空头持仓状态
            self.intra_trade_low = min(self.intra_trade_low, bar.low_price)
            self.intra_trade_high = bar.high_price

            # 检查卖出信号是否正确
            if self.last_signal_price > 0 and bar.close_price < self.last_signal_price:
                self.correct_signals += 1
                print(f"卖出信号正确: 价格从 {self.last_signal_price:.2f} 下跌到 {bar.close_price:.2f}")
                self.last_signal_price = 0.0

            # 设置止损点
            short_stop = self.intra_trade_low * (1 + self.trailing_percent / 100)
            self.cover(short_stop, abs(self.pos), stop=True)
            print(f"空头持仓，当前价格: {bar.close_price:.2f}, 止损价格: {short_stop:.2f}")

        self.put_event()
        
    def write_log(self, msg: str):
        """输出日志"""
        if self.cta_engine:
            self.cta_engine.write_log(msg, self.strategy_name) 
            
    def get_signal(self) -> Dict:
        """获取交易信号"""
        return {
            "buy": self.buy_signal,
            "sell": self.sell_signal
        }

    def get_parameters(self) -> Dict:
        """获取策略参数"""
        return self.parameters.copy()

    def set_parameters(self, params: Dict):
        """设置策略参数"""
        self.parameters.update(params)
        # 更新类属性中的参数
        for key, value in params.items():
            if hasattr(self, key):
                setattr(self, key, value)
